import dotenv
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from pydantic.v1 import BaseModel, Field
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()


# 1.创建一个json数据结构，用于高速大语言模型这个json长什么样
class Joke(BaseModel):
    # 冷笑话
    joke: str = Field(description="回答用户的冷笑话")

    # 冷笑话的笑点
    punchline: str = Field(description="回答用户的冷笑话的笑点")

parser = JsonOutputParser(pydantic_object=Joke)

# 2.构建一个提示模板
prompt = (ChatPromptTemplate.from_template("请根据用户的提问进行回答。 \n{format_instaructions} \n {query}")
          .partial(format_instaructions=parser.get_format_instructions()))

# 3.构建大语言模型
llm = ChatOpenAI(model_name="kimi-k2-0711-preview")

# 4.传递提示并进行解析
joke = parser.invoke(llm.invoke(prompt.format(query="请给我讲一个冷笑话")))
print(joke)
